package edu.arizona.sista.learn.language;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import start.Entity;
import start.Generator;
import start.RoleFrame;

import com.martiansoftware.jsap.FlaggedOption;
import com.martiansoftware.jsap.JSAP;
import com.martiansoftware.jsap.JSAPResult;
import com.martiansoftware.jsap.Switch;

import edu.arizona.sista.learn.activity.model.ActivityInstance;
import edu.arizona.sista.learn.activity.model.PredicateInterval;
import edu.arizona.sista.learn.model.ParamInterval;
import edu.arizona.sista.learn.util.Utils;

public class LanguageGeneration {
	
	private static String outDir = "data/" + Utils.OUTDIR + "sentences/";

	public static void main(String args[]) throws Exception {

		JSAP jsap = new JSAP();
		Switch helpSwitch = new Switch("help")
				.setShortFlag('h')
				.setLongFlag("help");
		FlaggedOption fluentsFlag = new FlaggedOption("fluents")
				.setDefault("")
				.setRequired(true)
				.setLongFlag("fluents");
		FlaggedOption out = new FlaggedOption("out")
				.setDefault("")
				.setRequired(false)
				.setLongFlag("out");
		FlaggedOption frames = new FlaggedOption("total-frames")
				.setDefault("")
				.setRequired(false)
				.setLongFlag("total-frames");
		jsap.registerParameter(helpSwitch);
		jsap.registerParameter(fluentsFlag);
		jsap.registerParameter(out);
		jsap.registerParameter(frames);

		JSAPResult config = jsap.parse(args);
		if (config.getBoolean("help") || "".equals(config.getString("fluents"))) {
			System.out.print("java -jar language-generation.jar ");
			System.out.println(jsap.getUsage());
			return;
		}
		
		File fluentsFile = new File(config.getString("fluents"));
		outDir = ("".equals(config.getString("out"))) ? outDir : config.getString("out");
		if (!outDir.endsWith("/")) outDir += "/";
		int totalFrames = ("".equals(config.getString("total-frames"))) ? 0
				: Integer.parseInt(config.getString("total-frames"));
		
		// Hard-coded for fast debug
//		fluentsFile = new File("/Users/trananh/Desktop/psi-y2/data-Jim/CR1_00002/fluents-cleaned/CR1_00002.lisp");
//		totalFrames = 1245;			/* map=420,newspaper=168,STOP=1213,
//		 							 * Y105=411,Y106=960,Y201=958,Y204=905,Y207=1287
//		 							 */
		
		// Load activity instances
		List<ActivityInstance> instances = ActivityInstance.load(fluentsFile);
		
		// Generate language and output
		List<ParamInterval<String>> sentenceIntervals = new ArrayList<ParamInterval<String>>();
		for (ActivityInstance inst: instances) {
			sentenceIntervals.addAll(generateLanguage(inst, outDir 
					+ fluentsFile.getName().replace(".lisp", ".txt")));
		}
		
		
		// Generate binary-vector output for sentences
		if (totalFrames > 0) {
			Map<String, int[]> sentenceMap = new HashMap<String, int[]>();
			List<String> sentences = new ArrayList<String>();
			for (ParamInterval<String> sentenceInterval : sentenceIntervals) {
				String s = sentenceInterval.obj();
				if (!sentenceMap.containsKey(s)) {
					sentences.add(s);
					sentenceMap.put(s, new int[totalFrames]);	// vector initialized to 0's
				}
				turnOnInterval(sentenceMap.get(s),
						sentenceInterval.interval().start(),
						sentenceInterval.interval().end());
			}
			
			// Output bit vectors
			outputSentenceBitVectors(outDir 
						+ fluentsFile.getName().replace(".lisp", ""),
						sentenceMap, sentences, totalFrames);
		}
	}
	

	private static List<ParamInterval<String>> generateLanguage(ActivityInstance inst, String outFile) {
		List<ParamInterval<String>> sentences = new ArrayList<ParamInterval<String>>();
		try {
			FileWriter fstream = new FileWriter(new File(outFile));
			BufferedWriter out = new BufferedWriter(fstream);
		
			// Grab the generator
			Generator generator = Generator.getGenerator();
		
			// Get all the intervals and sort by earliest starting time
			List<PredicateInterval> intervals = new ArrayList<PredicateInterval>();
			for (String key : inst.getIntervalsMap().keySet()) {
				intervals.addAll(inst.getIntervals(key));
			}
			Collections.sort(intervals, PredicateInterval.esf);
			
			// Produce a sentence for each predicate interval
			for (PredicateInterval pi : intervals) {
				String prefix = "["+pi.start+" "+pi.end+")" + "\t=> "; 
				String sentence = "";
				if (pi.getPredicate().getName().startsWith("pair")) {
					// stats-group(X,a,b,c) => A, B, and C are together with X percent confidence.
					String temp = pi.getPredicate().getArg(1).getBoundValue();
					if (pi.getPredicate().getArgs().length == 3) {
						temp =  temp + " and " + pi.getPredicate().getArg(2).getBoundValue();
					} else {
						for (int i = 2; i < pi.getPredicate().getArgs().length-1; i++) {
							temp = temp + ", " + pi.getPredicate().getArg(i).getBoundValue();;
						}
						temp = temp + ", and " +
							pi.getPredicate().getArg(pi.getPredicate().getArgs().length-1).getBoundValue();
					}
					temp = temp + " are together with at least " +
						pi.getPredicate().getArg(0).getBoundValue() +
						"% confidence.";
					sentence += temp;
				} else if (pi.getPredicate().getName().endsWith("group-formation")) {
					// group-formation(G,A,B,C,...,Z) => A is in group G with B, C, ..., and Z.
					assert pi.getPredicate().getArgs().length >= 3 : "group-formation must have at least 3 args";
					String temp = pi.getPredicate().getArg(1).getBoundValue();
					if (pi.getPredicate().getName().startsWith("social")) {
						temp += " is in social group ";
					} else if (pi.getPredicate().getName().startsWith("spt")) {
						temp += " is in spatio-temporal group ";
					} else {
						temp += " is in " 
							+ pi.getPredicate().getName().replace("-group-formation", "")
							+ " group ";
					}
					temp = temp + pi.getPredicate().getArg(0).getBoundValue() + " with ";
					if (pi.getPredicate().getArgs().length == 3) {
						temp = temp + pi.getPredicate().getArg(2).getBoundValue() + ".";
					} else {
						for (int i = 2; i < pi.getPredicate().getArgs().length-1; i++) {
							temp = temp + pi.getPredicate().getArg(i).getBoundValue() + ", ";
						}
						temp = temp + "and "
							+ pi.getPredicate().getArg(pi.getPredicate().getArgs().length-1).getBoundValue()
							+ ".";
					}
					sentence += temp;
				} else if (pi.getPredicate().getName().equalsIgnoreCase("holding")) {
					// holding(a, b, c) => A is holding C with hand-B.
					String body = pi.getPredicate().getArg(1).getBoundValue();
					if (pi.getPredicate().getArg(1).getBoundValue().contains(".")) {
						body = body.substring(body.indexOf('.')+1);
					}
					String temp = pi.getPredicate().getArg(0).getBoundValue()
						+ " is holding "
						+ pi.getPredicate().getArg(2).getBoundValue()
						+ " using "
						+ body
						+ ".";
					sentence += temp;
				} else {
					sentence += generator.generate(getRoleFrame(pi));
					if (sentence.substring(sentence.length()-1).equals("\n")) {
						sentence = sentence.substring(0, sentence.length()-1);
					}
				}
				
				if (!sentence.equalsIgnoreCase("null")){
					ParamInterval<String> sentenceInterval = new ParamInterval<String>(pi.start, pi.end, sentence);
					sentences.add(sentenceInterval);
					out.write(prefix + sentence + "\n");
				} else {
					out.write(prefix + "[ Unable to translate: " + pi.getPredicate().toString() + " ]\n");
				}
			}

			out.close();
		} catch (Exception e) {
			System.err.println("Error: " + e.getMessage());
		}
		
		return sentences;
	}
	
	private static RoleFrame getRoleFrame(PredicateInterval pi) {
		//Gestures {carry,catch,dig,fall,jump,kick,pickup,putdown,raise,run,throw,turn,walk}
		
		if (pi.getPredicate().getName().equalsIgnoreCase("movepos") ||
				pi.getPredicate().getName().equalsIgnoreCase("walk-gesture") ||
				pi.getPredicate().getName().equalsIgnoreCase("run-gesture")) {
			// movepos(a) => A is moving.
			String verb = pi.getPredicate().getName();
			if (verb.equals("movepos")) verb = "move";
			if (verb.equals("walk-gesture")) verb = "walk";
			if (verb.equals("run-gesture")) verb = "run";
			return new RoleFrame(new Entity(pi.getPredicate().getArg(0).getBoundValue()).noDeterminer(),
					verb).progressive();
		} else if (pi.getPredicate().getName().equalsIgnoreCase("carry") ||
				pi.getPredicate().getName().equalsIgnoreCase("catch") ||
				pi.getPredicate().getName().equalsIgnoreCase("raise") ||
				pi.getPredicate().getName().equalsIgnoreCase("throw")) {
			// carry(A) => A is carrying something.
			String verb = pi.getPredicate().getName();
			return new RoleFrame(new Entity(pi.getPredicate().getArg(0).getBoundValue()).noDeterminer(),
					verb, "something").progressive();
		} else if (pi.getPredicate().getName().equalsIgnoreCase("go")) {
			// go(a) => A started moving.
			return new RoleFrame(new Entity(pi.getPredicate().getArg(0).getBoundValue()).noDeterminer(),
					"start", "moving").past();
		} else if (pi.getPredicate().getName().equalsIgnoreCase("stop")) {
			// stop(a) => A stopped moving.
			return new RoleFrame(new Entity(pi.getPredicate().getArg(0).getBoundValue()).noDeterminer(),
					"stop", "moving").past();
		} else if (pi.getPredicate().getName().equalsIgnoreCase("approach") ||
				pi.getPredicate().getName().equalsIgnoreCase("chase") ||
				pi.getPredicate().getName().equalsIgnoreCase("follow")) {
			// approach(a, b) => B is approached by A  || A is approaching B.
//			if (new Random().nextInt(2) == 0)
//				return new RoleFrame(new Entity(pi.getPredicate().getArg(0).getBoundValue()).noDeterminer(),
//						pi.getPredicate().getName(),
//						new Entity(pi.getPredicate().getArg(1).getBoundValue()).noDeterminer())
//						.passive();
//			else
				return new RoleFrame(new Entity(pi.getPredicate().getArg(0).getBoundValue()).noDeterminer(),
						pi.getPredicate().getName(),
						new Entity(pi.getPredicate().getArg(1).getBoundValue()).noDeterminer())
						.progressive();
		} else if (pi.getPredicate().getName().equalsIgnoreCase("pass")) {
			// pass(a, b) => A passed by B.
			return new RoleFrame(new Entity(pi.getPredicate().getArg(0).getBoundValue()).noDeterminer(),
				"pass").addRole("by", new Entity(pi.getPredicate().getArg(1).getBoundValue()).noDeterminer())
				.past();
		} else if (pi.getPredicate().getName().equalsIgnoreCase("flee")) {
			// flee(a, b) => A fled from B quickly.
			return new RoleFrame(new Entity(pi.getPredicate().getArg(0).getBoundValue()).noDeterminer(),
				"flee").addRole("from", new Entity(pi.getPredicate().getArg(1).getBoundValue()).noDeterminer())
				.addModifier("quickly").past();
		} else if (pi.getPredicate().getName().equalsIgnoreCase("look-at")) {
			// look-at(a, b) => A is looking at B
			// look-at(a, b.body-part) => A is looking at B's arm.
			if (pi.getPredicate().getArg(1).getBoundValue().contains(".")) {
				String target = pi.getPredicate().getArg(1).getBoundValue();
				String targ = target.substring(0, target.indexOf('.'));
				String body = target.substring(target.indexOf('.')+1);
				return new RoleFrame(new Entity(pi.getPredicate().getArg(0).getBoundValue()).noDeterminer(),
						"look").addRole("at", new Entity(body).makePossessor(targ)).progressive();
			} else 
				return new RoleFrame(new Entity(pi.getPredicate().getArg(0).getBoundValue()).noDeterminer(),
					"look").addRole("at", new Entity(pi.getPredicate().getArg(1).getBoundValue()).noDeterminer())
					.progressive();
		} else if (pi.getPredicate().getName().equalsIgnoreCase("displaying")) {
			// displaying(A,B) => A is displaying B.
			return new RoleFrame(new Entity(pi.getPredicate().getArg(0).getBoundValue()).noDeterminer(),
				"display", new Entity(pi.getPredicate().getArg(1).getBoundValue()).noDeterminer())
				.progressive();
		} else if (pi.getPredicate().getName().equalsIgnoreCase("pickup") ||
				pi.getPredicate().getName().equalsIgnoreCase("pick-up")) {
			// pickup(A) => A is picking something up.
			return new RoleFrame(new Entity(pi.getPredicate().getArg(0).getBoundValue()).noDeterminer(),
					"pick", "something").addModifier("up").progressive();
		} else if (pi.getPredicate().getName().equalsIgnoreCase("putdown") ||
				pi.getPredicate().getName().equalsIgnoreCase("put-down")) {
			// putdown(A) => A is putting something down.
			return new RoleFrame(new Entity(pi.getPredicate().getArg(0).getBoundValue()).noDeterminer(),
					"put", "something").addModifier("down").progressive();
		} else if (pi.getPredicate().getArgs().length == 1) {
			// Default behavior for 1-parameter:
			// 	VERB(a) => A is "doing-verb".
			return new RoleFrame(new Entity(pi.getPredicate().getArg(0).getBoundValue()).noDeterminer(),
					pi.getPredicate().getName()).progressive();
		} else if (pi.getPredicate().getArgs().length == 2) {
			// Default behavior for 2-parameters:
			// 	VERB(a,b) => A is "doing-verb" B.
			return new RoleFrame(new Entity(pi.getPredicate().getArg(0).getBoundValue()).noDeterminer(),
				pi.getPredicate().getName(),
				new Entity(pi.getPredicate().getArg(1).getBoundValue()).noDeterminer())
				.progressive();
		}
		
		return new RoleFrame();
	}
	
	
	private static void outputSentenceBitVectors(String outprefix,
			Map<String, int[]> vectorsMap, List<String> sentences,
			int totalFrames) {
		try {
			// Write index file
			BufferedWriter out1 = new BufferedWriter(new FileWriter(new File(outprefix+"-idx.csv")));
			for (int i = 0; i < sentences.size(); i++) {
				out1.write("\""+sentences.get(i) + "\"\n");
			}
			out1.close();
				
			BufferedWriter out2 = new BufferedWriter(new FileWriter(new File(outprefix+"-bits.csv")));
			for (int f = 0; f < totalFrames; f++) {
				for (int i = 0; i < sentences.size(); i++) {
					int[] vector = vectorsMap.get(sentences.get(i));
					if (i == 0)
						out2.write(""+vector[f]);
					else
						out2.write(" "+vector[f]);
				}
				out2.write("\n");
			}
			
			out2.close();
		} catch (Exception e) {
			System.err.println("Error: " + e.getMessage());
		}
	}
	
	
	private static void turnOnInterval(int[] frames, int start, int end) {
		start = Math.max(start, 0);
		end = Math.min(end, frames.length);
		for (int i = start; i < end; i++)
			frames[i] = 1;
	}
	
	
	@SuppressWarnings("unused")
	private static void genesisDemo() {
		System.out.println("Running Start generator on test sentences");

		// Obtain singleton of generator class
		Generator generator = Generator.getGenerator();

		// Entities
		String test = "Person0.arm";
		String targ = test.substring(0, test.indexOf('.'));
		String body = test.substring(test.indexOf('.')+1);
		System.out.println(generator.generate(new RoleFrame("person0", "look").addRole("at", new Entity(body).makePossessor(targ)).progressive()));
		System.out.println(generator.generate(new RoleFrame(new Entity("men")
				.noDeterminer(), "run")));
	}
}
